package ws

import (
	"errors"
	"golang.org/x/net/websocket"
	"io/ioutil"
)

const (
	WS_MESSAGE_TEXT   = 1
	WS_MESSAGE_BINARY = 2
)

func ReceiveMessage(conn *websocket.Conn) (msgtype byte, msg []byte, err error) {
again:
	frame, err := conn.NewFrameReader()
	if err != nil {
		return 0, nil, err
	}
	frame, err = conn.HandleFrame(frame)
	if err != nil {
		return 0, nil, err
	}
	if frame == nil {
		goto again
	}
	maxPayloadBytes := conn.MaxPayloadBytes
	if maxPayloadBytes == 0 {
		maxPayloadBytes = websocket.DefaultMaxPayloadBytes
	}
	payloadType := frame.PayloadType()
	data, err := ioutil.ReadAll(frame)
	if err != nil {
		return 0, nil, err
	}
	return payloadType, data, err
}

func SendWsMessage(conn *websocket.Conn, msgtype byte, msg []byte) error {
	if conn == nil {
		return errors.New("尚未建立连接")
	}
	writer, err := conn.NewFrameWriter(msgtype)
	if err != nil {
		return err
	}
	_, err = writer.Write(msg)
	return err
}
